import cv2
import numpy as np
import os
import matplotlib.pyplot as plt
from pathlib import Path
import time

class PositionVisualizer:

    def __init__(self, rect_size=(0.8, 2.5), scale_factor=100):
        """
        初始化PositionVisualizer类
        
        Args:
            seg_model_path: 分割模型路径
            pose_model_path: 姿态模型路径
            device: 设备ID
            rect_size: 矩形大小（宽, 长），单位为米，默认为(0.8, 2.5)
            scale_factor: 米到像素的转换因子，默认为100（1米=100像素）
        """
        # self.yolo_seg_pose = YoloSegPose(seg_model_path, pose_model_path)
        # self.standard_vector = self.yolo_seg_pose.set_standard_position(standard_image_path)
        self.rect_size = rect_size  # (宽度0.8米, 长度2.5米)
        self.scale_factor = scale_factor  # 米到像素的转换因子
        self.output_dir = "results_position_visualization"
        
        # 创建输出目录
        os.makedirs(self.output_dir, exist_ok=True)
    
    def draw_rectangle(self, vector, ax, label_prefix="Standard", color=(0, 1, 0), thickness=2):
        """
        使用Matplotlib在指定轴上绘制矩形、向量和角点，矩形固定为2.5米x0.8米
        
        Args:
            vector: 位置向量，包含以下结构：
                - center: 转换后的中心点坐标（矩形左下角点，米）
                - bottom_vec: 转换后的x方向向量（米）
                - left_vec: 转换后的y方向向量（米）
                - z_vec: 转换后的z方向向量（不使用）
                - bottom_angle: bottom向量与x轴的夹角（不使用）
                - left_angle: left向量与y轴的夹角（不使用）
            ax: Matplotlib轴对象
            label_prefix: 标签前缀（标准或测试）
            color: 矩形颜色 (RGB格式, 默认绿色，范围[0,1])
            thickness: 矩形线条粗细
        """
        if vector is None:
            return
        # 从向量中提取中心点坐标和向量（单位为米）
        # print("vector:",vector)
        center_x, center_y = vector['center'][0], vector['center'][1]
        left_vec = vector['left_vec'][:2]
        bottom_vec = vector['bottom_vec'][:2]
        
        # 归一化向量并缩放到固定尺寸（宽0.8米，长2.5米）
        rect_width, rect_height = self.rect_size  # 宽0.8米，长2.5米
        lx, ly = left_vec
        len_left = (lx**2 + ly**2)**0.5
        bx, by = bottom_vec
        len_bottom = (bx**2 + by**2)**0.5
        
        # 归一化并缩放left_vec到2.5米（长度）
        if len_left > 1e-5:
            left_vec = (lx / len_left * rect_height, ly / len_left * rect_height)
        else:
            left_vec = (0, rect_height)
        # 归一化并缩放bottom_vec到0.8米（宽度）
        if len_bottom > 1e-5:
            bottom_vec = (bx / len_bottom * rect_width, by / len_bottom * rect_width)
        else:
            bottom_vec = (rect_width, 0)
        
        # 应用像素缩放（1米 = scale_factor像素）
        center_x, center_y = center_x * self.scale_factor, center_y * self.scale_factor
        left_vec = (left_vec[0] * self.scale_factor, left_vec[1] * self.scale_factor)
        bottom_vec = (bottom_vec[0] * self.scale_factor, bottom_vec[1] * self.scale_factor)
        
        # 计算向量长度（像素）
        lx, ly = left_vec
        len_left = (lx**2 + ly**2)**0.5
        bx, by = bottom_vec
        len_bottom = (bx**2 + by**2)**0.5
        
        center = (center_x, center_y)  # 矩形左下角点（像素）
        # 计算与left_vec垂直的向量
        if len_left > 1e-5:
            perp1 = (ly / len_left, -lx / len_left)  # 顺时针旋转90°
            perp2 = (-ly / len_left, lx / len_left)  # 逆时针旋转90°
            dot1 = bx * perp1[0] + by * perp1[1]
            dot2 = bx * perp2[0] + by * perp2[1]
            if dot1 > dot2:
                adjusted_bottom_vec = (perp1[0] * len_bottom, perp1[1] * len_bottom)
            else:
                adjusted_bottom_vec = (perp2[0] * len_bottom, perp2[1] * len_bottom)
        else:
            adjusted_bottom_vec = bottom_vec

        # 计算矩形的四个角点（像素）
        bottom_left = center
        bottom_right = (center_x + adjusted_bottom_vec[0], center_y + adjusted_bottom_vec[1])
        top_left = (center_x + left_vec[0], center_y + left_vec[1])
        top_right = (top_left[0] + adjusted_bottom_vec[0], top_left[1] + adjusted_bottom_vec[1])
        
        # 调试：打印角点坐标
        # print(f"{label_prefix}矩形角点: 左下={bottom_left}, 右下={bottom_right}, 左上={top_left}, 右上={top_right}")
        
        # 绘制矩形
        from matplotlib.patches import Polygon
        pts = np.array([bottom_left, bottom_right, top_right, top_left])
        polygon = Polygon(pts, closed=True, edgecolor=color, linewidth=thickness, fill=False, label=f"{label_prefix}Rect")
        ax.add_patch(polygon)
        
        
    def visualize_position(self,standard_vector, test_vector,timestamp):
        try:
            # visualizer = PositionVisualizer(standard_vector,test_vector, rect_size, scale_factor)
            # 设置Matplotlib交互模式
            plt.ion()  # 开启交互模式
            fig, ax = plt.subplots(figsize=(8, 8))  # 创建单一图形窗口
            # 设置坐标轴范围
            rect_width, rect_height = self.rect_size
            max_vec_length = max(rect_width, rect_height) * self.scale_factor 
            canvas_width = int(max_vec_length + 200)
            canvas_height = int(max_vec_length + 200)
        
            # 清空当前轴内容
            ax.clear()
            
            if test_vector is None:
                print("无法计算位置向量，跳过此图像")
                return
            # print(f"测试位置向量: {test_vector}")
            ax.set_title(f"Standard (440.png) vs Test {timestamp}", fontsize=12)

            # 绘制标准和测试矩形
            self.draw_rectangle(standard_vector, ax, label_prefix="Standard", color=(0, 1, 0), thickness=2)
            self.draw_rectangle(test_vector, ax, label_prefix="Test", color=(0, 0, 1), thickness=2)
            
            # 设置坐标轴范围
            ax.set_xlim(-50, canvas_width)
            ax.set_ylim(-50, canvas_height)
            ax.set_aspect('equal')
            ax.legend()
            
            # 刷新显示
            # plt.draw()
            # plt.pause(0.1)
            
            # 保存结果
            output_path = os.path.join(self.output_dir, f"position_{timestamp}")
            fig.savefig(output_path, dpi=100, bbox_inches='tight')
            print(f"结果已保存至: {output_path}")
            
                    
        except Exception as e:
            import traceback
            print(f"执行过程中发生错误: {e}")
            print(traceback.format_exc())
            
        finally:
            # 确保图形窗口关闭
            plt.ioff()
            plt.close('all')
            print("程序结束")

